package com.lovely602.ai.alibaba.controller.embedding;


import com.alibaba.cloud.ai.dashscope.embedding.DashScopeEmbeddingOptions;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.embedding.EmbeddingRequest;
import org.springframework.ai.embedding.EmbeddingResponse;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;

import java.util.List;
import java.util.Map;

/**
 *
 * @author lizhixing
 */
@Slf4j
@RestController
@RequestMapping(value = "/embedding")
public class EmbeddingController {

    private final EmbeddingModel embeddingModel;

    public EmbeddingController(EmbeddingModel embeddingModel) {
        this.embeddingModel = embeddingModel;
    }

    /**
     * 将文本转换成向量
     */
    @GetMapping("/embed")
    public Map embed(@RequestParam(value = "msg", defaultValue = "Tell me a joke") String msg) {
        EmbeddingResponse embeddingResponse = this.embeddingModel.embedForResponse(List.of(msg));
        return Map.of("embedding", embeddingResponse);
    }

    /**
     * 自定义嵌入模型参数
     */
    @GetMapping("/custom")
    public Map custom(@RequestParam(value = "msg", defaultValue = "Tell me a joke") String msg) {

        // 嵌入模型参数
        DashScopeEmbeddingOptions embeddingOptions = DashScopeEmbeddingOptions.builder()
                .withModel("text-embedding-v2")
                .build();
        // 嵌入请求
        EmbeddingRequest embeddingRequest = new EmbeddingRequest(List.of(msg), embeddingOptions);

        EmbeddingResponse embeddingResponse = embeddingModel.call(embeddingRequest);

        return Map.of("embedding", embeddingResponse);
    }

}
